from base64 import b64encode
import requests
from ast import literal_eval

def get_default():
    "调用本地7861端口的/接口，用来测试能否正常连接服务"
    res = requests.get('http://127.0.0.1:7861/')
    if res.status_code == 200:
        return "LangChain-ChatChat服务连接成功！"
    else:
        return 'LangChain-ChatChat服务连接失败！'

def bytes2base64(b:bytes) ->bytes:
    return b64encode(b).decode('utf8')

def chatchat(input:str="我们来玩成语接龙，我先来，生龙活虎",model="Qwen-0.5B-Chat"):
    "调用本地7861端口的/chat/chat接口，直接与大模型对话(非流式输出)"
    # 组织输入
    xjson = {
  "query": input,
  "conversation_id": "",
  "stream": False,
  "model_name": model,
  "temperature": 0.7,
  "max_tokens": 0,
  "prompt_name": "default"
}

    # 获取输出
    res = requests.post('http://127.0.0.1:7861/chat/chat/', json=xjson)
    res = literal_eval(res.text[6:])

    # 返回结果
    return res["text"]

def knowledge_base_chat(kb='Demo',input="如何使用git?",model="Qwen-0.5B-Chat"):
    """ 使用本地Demo知识库，输入："如何使用git？" """
    xjson = {
  "query": input,
  "knowledge_base_name": kb,
  "top_k": 3,
  "score_threshold": 1,
  "stream": False,
  "model_name": model,
  "temperature": 0.7,
  "max_tokens": 0,
  "prompt_name": "default"
}
    
    res = requests.post('http://127.0.0.1:7861/chat/knowledge_base_chat/', json=xjson)
    # res = literal_eval(res.text[141:])
    res = literal_eval(res.text.split('data: ')[1][:-4])
    answer = res['answer']
    docs = res['docs']
    return answer, docs

# def file_chat(query="张三的儿子是谁？",file_path=r"E:\haymaker\AIdata\BW\ZG6\Projects\test_file_chat.txt"):
#     #开发中，未完成
#     with open(file_path, 'rb') as f:
#         file_bytes = f.read()
#     file = bytes2base64(file_bytes)
#     xjson = {
#   "query": query,
#   "knowledge_id": file,
#   "top_k": 3,
#   "score_threshold": 1,
#   "history": [
#   ],
#   "stream": False,
#   "model_name": "Qwen-0.5B-Chat",
#   "temperature": 0.7,
#   "max_tokens": 0,
#   "prompt_name": "default"
# }
#     res = requests.post('http://127.0.0.1:7861/chat/file_chat/', json=xjson)
#     return res.text

def list_knowledge_bases():
    "获取知识库列表"
    res = requests.get('http://127.0.0.1:7861/knowledge_base/list_knowledge_bases/')
    res = res.text
    return literal_eval(res)['data']

def list_files(kb='Demo'):
    "获取知识库kb内的文件列表"
    res = requests.get(f'http://127.0.0.1:7861/knowledge_base/list_files?knowledge_base_name={kb}')
    return literal_eval(res.text)['data']

if '__main__' == __name__:

    # 调用封装的函数
    print(get_default())
    print('chatchat output:', chatchat())
    answer,docs = knowledge_base_chat()
    print(f'\n大模型知识库问答结果:{answer}\n\n参考文档:\n')
    for doc in docs:
        print(doc)
    print(list_files())

    print(list_knowledge_bases())